"""
=========================================================
通用 Grover 算子构建模块（基于 cqlib）
=========================================================

功能：
------
该模块用于在 cqlib 框架下自动生成 Grover 算子：
    Q = A · S₀ · A† · S_f
其中：
    - A        : 状态制备电路（把 |0…0⟩ 变成初始叠加态 |ψ₀⟩）
    - S₀       : 对 |0…0⟩ 态的反射（diffusion about zero）
    - S_f      : Oracle（对“目标态”加上 -1 相位）
    - Q        : Grover 扩散算子（一次“放大迭代”）

该模块自动处理：
    - Oracle 组合；
    - 状态制备及其逆；
    - 多控-Z 分解；
    - X → 多控Z → X 的零反射；
    - 支持任意数量控制位（带或不带 ancilla）；
    - 允许 barrier 插入（调试可视化）。

依赖：
    - circuits.algorithm_circuit.Algorithm_Circuit
    - cqlib.circuits.qubit.Qubit

"""

from typing import List, Optional, Tuple
import numpy as np

from cqlib.circuits.qubit import Qubit
from circuits.algorithm_circuit import Algorithm_Circuit as Circuit


# ============================================================
# 多控制 X 门（MCX）的链式分解
# ============================================================
def _multi_controlled_x_with_ancillas(
    circ: Circuit,
    controls: List[int],
    target: int,
    ancillas: List[int],
):
    """
    实现任意控制位数量的多控 X 门 (MCX)。

    原理：
        - 若控制位 <= 2：直接用 CX / CCX；
        - 若控制位 > 2：构造一条 Toffoli 链，
          使用 (k-2) 个 ancilla 累积前面控制位的与操作。

    链式逻辑：
        a0 = c0 & c1
        a1 = c2 & a0
        a2 = c3 & a1
        ...
        最后：ccx(c_{k-1}, a_{k-3}, target)

    用完后需反向撤销 ancilla 的计算。

    注意：
        需要至少 (k-2) 个 ancilla；否则抛出异常。
    """
    k = len(controls)
    if k == 1:
        circ.cx(controls[0], target)
        return
    if k == 2:
        circ.ccx(controls[0], controls[1], target)
        return

    need = k - 2
    if len(ancillas) < need:
        raise ValueError(
            f"多控X需要 {need} 个 ancilla，但仅提供 {len(ancillas)} 个。"
            " 请在反射集之外预留足够的辅比特，或改用相对相位分解。"
        )

    # === 前向累积阶段 ===
    circ.ccx(controls[0], controls[1], ancillas[0])
    for i in range(2, k - 1):
        circ.ccx(controls[i], ancillas[i - 2], ancillas[i - 1])

    # === 作用在目标位 ===
    last_a = ancillas[need - 1]  # 最后一个累积位
    circ.ccx(controls[k - 1], last_a, target)

    # === 反向撤销 ancilla 状态 ===
    for i in reversed(range(2, k - 1)):
        circ.ccx(controls[i], ancillas[i - 2], ancillas[i - 1])
    circ.ccx(controls[0], controls[1], ancillas[0])


# ============================================================
# 多控 Z 的实现：通过 H-变换转为 MCX
# ============================================================
def _ncz_from_mcx(circ: Circuit, qubits: List[int], ancillas: List[int]):
    """
    把多控 Z 门实现为：
        H(target) → 多控 X → H(target)
    即：将目标 Z 门转化为 X 门控制链（利用 HZH=X）。
    """
    controls, target = qubits[:-1], qubits[-1]
    circ.h(target)
    _multi_controlled_x_with_ancillas(circ, controls, target, ancillas)
    circ.h(target)


# ============================================================
# 主函数：Grover 算子构造
# ============================================================
def grover_operator(
    oracle: Circuit,
    state_preparation: Optional[Circuit] = None,
    reflection_qubits: Optional[List[int]] = None,
    insert_barriers: bool = False,
    name: str = "Q",
) -> Circuit:
    """
    构造通用 Grover 算子：
        Q = A · S₀ · A† · S_f

    参数：
        oracle             : Circuit，Oracle 电路（实现 S_f，即标记解的相位翻转）
        state_preparation  : Circuit，可选，状态制备电路 A；
                             若为 None，则默认 A = H^{⊗n}
        reflection_qubits  : 反射操作所作用的量子位索引；
                             若为 None，则默认对全体量子位。
        insert_barriers    : 是否在关键阶段插入 barrier 以便可视化；
        name               : 电路命名（可忽略）。

    流程：
        (1) Oracle (S_f)
        (2) A†（状态制备逆）
        (3) 零反射 S₀ = X → 多控Z → X
        (4) 再施加 A

    返回：
        构造完成的 Grover 算子电路 Circuit 对象。
    """
    nq = oracle.num_qubits
    circ = Circuit(nq)  # 创建一个新的空电路

    # ---------- (1) Oracle 操作 S_f ----------
    # Oracle 通常对“解”态加上相位 π（-1 因子）
    # 如果 cqlib 支持 compose，可直接合并
    circ = circ.compose(oracle)
    if insert_barriers:
        circ.barrier()

    # ---------- (2) 状态制备的逆 A† ----------
    if reflection_qubits is None:
        reflection_qubits = list(range(nq))

    if state_preparation is None:
        # 默认 A=H^{⊗n}，其逆等于自身
        for q in reflection_qubits:
            circ.h(q)
    else:
        # 若有自定义 A，则取其逆电路
        circ = circ.compose(state_preparation.inverse())

    if insert_barriers:
        circ.barrier()

    # ---------- (3) 零态反射 S₀ ----------
    # S₀ = X^{⊗n} Z₀ X^{⊗n}，其中 Z₀ 是 |11...1⟩ 上的相位翻转
    for q in reflection_qubits:
        circ.x(q)

    n_reflect = len(reflection_qubits)
    if n_reflect == 1:
        circ.z(reflection_qubits[0])
    else:
        target = reflection_qubits[-1]
        controls = reflection_qubits[:-1]

        # 自动选择 ancilla：从非反射位中抽取
        need_anc = max(0, len(controls) - 2)
        pool = [q for q in range(nq) if q not in reflection_qubits]
        ancillas = pool[:need_anc]

        if len(controls) == 1:
            # 只有一个控制位：直接 CZ
            if hasattr(circ, "cz"):
                circ.cz(controls[0], target)
            else:
                # 若无 CZ，退化为 H-X-H 构造
                circ.h(target)
                circ.cx(controls[0], target)
                circ.h(target)
        else:
            # 多控制情况：构造 NCZ（多控Z）
            _ncz_from_mcx(circ, reflection_qubits, ancillas)

    # 撤销前面的 X
    for q in reflection_qubits:
        circ.x(q)

    if insert_barriers:
        circ.barrier()

    # ---------- (4) 再施加状态制备 A ----------
    if state_preparation is None:
        for q in reflection_qubits:
            circ.h(q)
    else:
        circ = circ.compose(state_preparation)

    return circ
